package play;

import org.antlr.v4.runtime.ParserRuleContext;

import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

/**
 * 基于 AST，产生汇编代码。
 * 1。仅支持整型，字符串仅支持字面量。
 * 2。运算仅支持加减法。
 */

public class ASMGen extends PlayScriptBaseVisitor<String> {

    // 之前的编译结果
    private AnnotatedTree at = null;

    public ASMGen(AnnotatedTree at) {
        this.at = at;
    }

    // 一些临时变量,用于保存扫描过程中的累计量
    // rsp的偏移量，栈桢的大小
    private int rspOffset = 0;

    // 计算过程中的临时变量存放的位置
    // key: AST 节点
    // value: 该节点在栈中存储的地址
    private Map<ParserRuleContext, String> tempVars = new HashMap<>();

    // 本地变量存储地址。目前只在栈中分配，不优化到寄存器中。
    // key: 变量
    // value：该节点存储的地址
    private Map<Variable, String> localVars = new HashMap<>();

    // 可以使用的寄存器
    private String[] registersl = {"%eax", "%ebx",
            "%r10d", "%r11d", "%r12d", "%r13d", "%r14d", "%r15d"};
    private String[] paramRegisterl = {"%edi", "%esi", "%edx", "%ecx", "%r8d", "%r9d"};
    private String[] paramRegisterq = {"%rdi", "%rsi", "%rdx", "%rcx", "%r8", "%r9"};

    // 扫描某个例程时,生成的代码缓存区
    private StringBuilder bodyASM = new StringBuilder();

    // 字符串字面量
    private List<String> stringLiterals = new LinkedList<>();


    // 主程序
    public String generate() {
        StringBuilder sb = new StringBuilder();
        // 1.代码段头
        sb.append("\t.section   __TEXT, __text, regular, pure_instructions\n");

        // 2.生成函数的代码
        for (Type type : at.types) {
            if (type instanceof Function) {
                Function function = (Function) type;
                PlayScriptParser.FunctionDeclarationContext fdc = (PlayScriptParser.FunctionDeclarationContext) function.ctx;
                // 遍历,生成代码到 bodyASM 中
                visitFunctionDeclaration(fdc);
                generateProcedure(function.name, sb);
            }
        }
        // 3.对主程序生成 main函数
        visitProg((PlayScriptParser.ProgContext) at.ast);
        generateProcedure("main", sb);

        // 4.文本字面量
        sb.append("\n# 字符串字面量\n");
        sb.append("\t.section  __TEXT, __cstring, cstring_literals\n");
        for (int i = 0; i < stringLiterals.size(); i++) {
            sb.append("L.str." + i + ":\n");
            sb.append("\t.asciz\t\"").append(stringLiterals.get(i)).append("\"\n");
        }

        // 5.重置全局临时变量
        stringLiterals.clear();
        return sb.toString();
    }

    /**
     * 生成过程体
     *
     * @param name 名称
     * @param sb   生成的 builder
     */
    private void generateProcedure(String name, StringBuilder sb) {
        // 1. 函数标签
        sb.append("\n## 过程:").append(name).append("\n");
        sb.append("\t.globl _").append(name).append("\n");
        sb.append("_").append(name).append(":\n");

        // 2. 序曲
        sb.append("\n\t# 序曲\n");
        sb.append("\tpushq\t%rbp\n");
        sb.append("\tmovq\t%rsp, %rbp\n");

        // 3.设置栈顶
        if ((rspOffset % 16) != 0) {
            // 16字节对齐
            rspOffset = (rspOffset / 16 + 1) * 16;
        }
        sb.append("\n\t# 设置栈顶\n");
        sb.append("\tsubq\t$").append(rspOffset).append(", %rsp\n");

        // 4.保存用到的寄存器
        saveRegisters();

        // 5.函数体
        sb.append("\n\t# 过程体\n");
        sb.append(bodyASM);

        // 6.恢复保护的寄存器
        restoreRegisters();

        // 7.恢复栈顶
        sb.append("\n\t# 恢复栈顶\n");
        sb.append("\taddq\t$").append(rspOffset).append(", %rsp\n");

        // 8.如果是main，设置返回值为0
        if (name.equals("main")) {
            sb.append("\n\t# 返回值\n");
            sb.append("\txorl\t%eax, %eax\n");
        }

        // 9.尾声
        sb.append("\n\t# 尾声\n");
        sb.append("\tpopq\t%rbp\n");
        sb.append("\tretq\n");

        // 10.重置临时变量
        rspOffset = 0;
        localVars.clear();
        tempVars.clear();
        bodyASM = new StringBuilder();
    }

    // 保存调用者需要保护的寄存器
    private void saveRegisters() {
        // TODO save registers
    }

    // 恢复被保护的寄存器
    private void restoreRegisters() {
        // TODO restore register
    }

    private String allocForExpression(PlayScriptParser.ExpressionContext ctx) {
        String rtn = null;

        // 复用前序表达式的存储位置
        if (ctx.bop != null && ctx.expression().size() >= 2) {
            PlayScriptParser.ExpressionContext left = ctx.expression(0);
            String leftAddress = tempVars.get(left);
            if (leftAddress != null) {
                // 当前节点也跟这个地址关联起来。
                tempVars.put(ctx, leftAddress);
                return leftAddress;
            }
        }

        int availableRegister = getAvailableRegister();
        if (availableRegister != -1) {
            rtn = registersl[availableRegister];
        } else {
            rspOffset += 4;
            rtn = "-" + rspOffset + "%rbp";
        }
        tempVars.put(ctx, rtn);
        return rtn;
    }

    private String getStringLiteralAddress(String string) {
        int index = stringLiterals.indexOf(string);
        if (index == -1) {
            stringLiterals.add(string);
            index = stringLiterals.size() - 1;
        }
        return "ref:L.str." + index + "(%rip)";
    }

    private int getAvailableRegister() {
        for (int i = 0; i < registersl.length; i++){
            String r = registersl[i];
            if (!tempVars.containsValue(r)) {
                return i;
            }
        }
        return -1;
    }

    // 继承 visitor 方法，用于产生 asm 代码片段

    @Override
    public String visitFunctionDeclaration(PlayScriptParser.FunctionDeclarationContext ctx) {
        // 给所有参数确定地址
        Function function = (Function) at.node2Scope.get(ctx);
        for (int i = 0; i < function.parameters.size(); i++) {
            if (i < 6) {
                // 少于6个参数，使用寄存器
                localVars.put(function.parameters.get(i), paramRegisterl[i]);
            } else {
                // 参数在栈中相对于 %rbp 的偏移量
                int paramOffset = (i - 6) * 8 + 16;
                String paramAddress = "" + paramOffset + "(%rbp)";
                localVars.put(function.parameters.get(i), paramAddress);
            }
        }
        return visitFunctionBody(ctx.functionBody());
    }

    @Override
    public String visitFunctionBody(PlayScriptParser.FunctionBodyContext ctx) {
        String value = null;
        if (ctx.block() != null) {
            value = visitBlock(ctx.block());
        }
        return value;
    }

    @Override
    public String visitBlock(PlayScriptParser.BlockContext ctx) {
        return visitBlockStatements(ctx.blockStatements());
    }

    @Override
    public String visitProg(PlayScriptParser.ProgContext ctx) {
        return visitBlockStatements(ctx.blockStatements());
    }

    @Override
    public String visitBlockStatements(PlayScriptParser.BlockStatementsContext ctx) {
        StringBuilder sb = new StringBuilder();
        for (PlayScriptParser.BlockStatementContext child : ctx.blockStatement()) {
            sb.append(visitBlockStatement(child));
        }
        return sb.toString();
    }

    @Override
    public String visitBlockStatement(PlayScriptParser.BlockStatementContext ctx) {
        StringBuilder sb = new StringBuilder();
        ;
        if (ctx.variableDeclarators() != null) {
            sb.append(visitVariableDeclarators(ctx.variableDeclarators()));
        } else if (ctx.statement() != null) {
            sb.append(visitStatement(ctx.statement()));
        }
        return sb.toString();
    }

    @Override
    public String visitStatement(PlayScriptParser.StatementContext ctx) {
        String value = "";
        if (ctx.statementExpression != null) {
            value = visitExpression(ctx.statementExpression);
        } else if (ctx.RETURN() != null) {
            if (ctx.expression() != null) {
                value = visitExpression(ctx.expression());
                bodyASM.append("\n\t# 返回值\n");
                if (value.equals("%eax")) {
                    bodyASM.append("\t# 返回值在之前的计算中，已经存入%eax\n");
                } else {
                    bodyASM.append("\tmovl\t").append(value).append(", %eax\n");
                }
            }
        }
        return value;
    }

    @Override
    public String visitVariableDeclarators(PlayScriptParser.VariableDeclaratorsContext ctx) {
        StringBuilder sb = new StringBuilder();
        for (PlayScriptParser.VariableDeclaratorContext child : ctx.variableDeclarator()) {
            sb.append(visitVariableDeclarator(child));
        }
        return sb.toString();
    }

    @Override
    public String visitVariableDeclarator(PlayScriptParser.VariableDeclaratorContext ctx) {
        String varAddress = visitVariableDeclaratorId(ctx.variableDeclaratorId());
        if (ctx.variableInitializer() != null) {
            String value = visitVariableInitializer(ctx.variableInitializer());
            bodyASM.append("\tmovl\t").append(value).append(", ").append(varAddress).append("\n");
        }
        return varAddress;
    }

    @Override
    public String visitVariableDeclaratorId(PlayScriptParser.VariableDeclaratorIdContext ctx) {
        // 整型变量 占4字节
        rspOffset += 4;
        String rtn = "-" + rspOffset + "(%rbp)";

        Symbol symbol = at.symbolOfNode.get(ctx);
        localVars.put((Variable) symbol, rtn);
        return rtn;
    }

    @Override
    public String visitVariableInitializer(PlayScriptParser.VariableInitializerContext ctx) {
        String rtn = "";
        if (ctx.expression() != null) {
            rtn = visitExpression(ctx.expression());
        }
        return rtn;
    }

    @Override
    public String visitExpression(PlayScriptParser.ExpressionContext ctx) {
        String address = "";
        if (ctx.bop != null && ctx.expression().size() >= 2) {
            // 二元运算
            String left = visitExpression(ctx.expression(0));
            String right = visitExpression(ctx.expression(1));
            switch (ctx.bop.getType()) {
                case PlayScriptParser.ADD:
                    // 为加法运算申请一个目标临时的存储位置。
                    address = allocForExpression(ctx);
                    if (!address.equals(left)) {
                        bodyASM.append("\tmovl\t").append(left).append(", ").append(address).append("\n");
                    }
                    bodyASM.append("\taddl\t").append(right).append(", ").append(address).append("\n");
                    break;
                case PlayScriptParser.SUB:
                    address = allocForExpression(ctx);
                    bodyASM.append("\tmovl\t").append(left).append(", ").append(address).append("\n");
                    bodyASM.append("\tsubl\t").append(right).append(", ").append(address).append("\n");
                    break;
                case PlayScriptParser.ASSIGN:
                    bodyASM.append("\tmovl\t").append(right).append(", ").append(left).append("\n");
                    break;
            }
        } else if (ctx.primary() != null) {
            address = visitPrimary(ctx.primary());
        } else if (ctx.functionCall() != null) {
            address = visitFunctionCall(ctx.functionCall());
        }
        return address;
    }

    @Override
    public String visitPrimary(PlayScriptParser.PrimaryContext ctx) {
        String rtn = "";
        if (ctx.literal() != null) {
            // 直接操作数
            rtn = visitLiteral(ctx.literal());
        } else if (ctx.IDENTIFIER() != null) {
            Symbol symbol = at.symbolOfNode.get(ctx);
            if (symbol instanceof Variable) {
                // 本地变量值，暂不支持上一级 Scope 的变量
                rtn = localVars.get((Variable) symbol);
            }
        }
        return rtn;
    }

    @Override
    public String visitLiteral(PlayScriptParser.LiteralContext ctx) {
        String rtn = "";
        if (ctx.integerLiteral() != null) {
            rtn = visitIntegerLiteral(ctx.integerLiteral());
        } else if (ctx.STRING_LITERAL() != null) {
            String withQuotationMark = ctx.STRING_LITERAL().getText();
            String withoutQuotationMark = withQuotationMark.substring(1, withQuotationMark.length() - 1);
            rtn = getStringLiteralAddress(withoutQuotationMark);
        }
        return rtn;
    }

    @Override
    public String visitIntegerLiteral(PlayScriptParser.IntegerLiteralContext ctx) {
        String rtn = "";
        if (ctx.DECIMAL_LITERAL() != null) {
            rtn = "$" + ctx.DECIMAL_LITERAL().getText();
        }
        return rtn;
    }

    @Override
    public String visitFunctionCall(PlayScriptParser.FunctionCallContext ctx) {
        // 默认返回值的地方
        String address = "%eax";
        String functionName = null;
        Symbol symbol = at.symbolOfNode.get(ctx);

        if (symbol instanceof Function) {
            Function function = (Function) symbol;
            functionName = function.name;
        } else {
            // TODO 临时打印输出
            if (ctx.IDENTIFIER().getText().equals("println")) {
                functionName = "printf";
            } else {
                at.log("unable to find function" + ctx.IDENTIFIER().getText(), ctx);
            }
        }

        // 设置参数
        if (ctx.expressionList() == null) {
            return address;
        }
        int paramOffset = 0;
        int numParams = ctx.expressionList().expression().size();
        // 1. 先计算所有参数的值。这个时候可能会引起栈的变化，用来存放临时变量
        int oldOffset = rspOffset;
        List<String> values = new LinkedList<>();
        for (int i = 0; i < numParams; i++) {
            values.add(visitExpression(ctx.expressionList().expression(i)));
        }
        int offset1 = rspOffset - oldOffset;
        // 2. 扩展栈
        if (numParams > 6) {
            paramOffset = 8 * (numParams - 6) + offset1;
            bodyASM.append("\n\t# 为参数扩展栈\n");
            bodyASM.append("\tsubq\t$").append(paramOffset).append(", %rsp\n");
        }
        // 3. 设置参数
        if (numParams > 0) {
            bodyASM.append("\n\t# 设置参数\n");
        }
        for (int i = 0; i < numParams; i++) {
            String value = values.get(i);
            String paramAddress = "";
            if (i < 6) {
                if (value.startsWith("ref:")) {
                    paramAddress = paramRegisterq[i];
                } else {
                    paramAddress = paramRegisterl[i];
                }
            } else {
                if (i == 6) {
                    paramAddress = "(%rsp)";
                } else {
                    paramAddress = "" + ((i - 6) * 8) + "(%rsp)";
                }
            }
            if (value.startsWith("ref:")) {
                // 传地址, lea = load effective address,
                // ref: https://stackoverflow.com/questions/1658294/whats-the-purpose-of-the-lea-instruction
                bodyASM.append("\tleaq\t").append(value.substring(4)).append(", ").append(paramAddress).append("\n");
            } else {
                // 传值
                bodyASM.append("\tmovl\t").append(value).append(", ").append(paramAddress).append("\n");
            }
        }

        // 4. 调用函数
        bodyASM.append("\n\t# 调用函数\n");
        bodyASM.append("\tcallq\t_").append(functionName).append("\n");

        // 5. 恢复栈
        if (numParams > 6) {
            paramOffset = 8 * (numParams - 6);
            bodyASM.append("\n\t# 收回参数的栈空间\n");
            bodyASM.append("\taddq\t$").append(paramOffset).append(", %rsp\n");
        }

        return address;
    }

    // 继承，但未直接使用

    @Override
    public String visitExpressionList(PlayScriptParser.ExpressionListContext ctx) {
        return super.visitExpressionList(ctx);
    }

    @Override
    public String visitFormalParameter(PlayScriptParser.FormalParameterContext ctx) {
        return super.visitFormalParameter(ctx);
    }

    @Override
    public String visitFormalParameterList(PlayScriptParser.FormalParameterListContext ctx) {
        return super.visitFormalParameterList(ctx);
    }

    @Override
    public String visitFormalParameters(PlayScriptParser.FormalParametersContext ctx) {
        return super.visitFormalParameters(ctx);
    }

    @Override
    public String visitPrimitiveType(PlayScriptParser.PrimitiveTypeContext ctx) {
        return super.visitPrimitiveType(ctx);
    }

    @Override
    public String visitTypeType(PlayScriptParser.TypeTypeContext ctx) {
        return super.visitTypeType(ctx);
    }
}
